#include "l1lossgrad_custom_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_l1lossgrad_custom.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void l1lossgrad_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z,GM_ADDR s,GM_ADDR reduction,GM_ADDR u,L1LossGradCustomTilingData tiling);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    size_t tilingSize = 2 * sizeof(uint32_t);
    size_t inputByteSize = 8* 2048 * sizeof(uint32_t);
    size_t outputByteSize = 8* 2048 *sizeof(uint32_t);
    size_t outputByteSize2 = 128 *sizeof(uint32_t);
    uint32_t reduction = 2;

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *y = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *u = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *z = (uint8_t *)AscendC::GmAlloc(outputByteSize);
    uint8_t *s = (uint8_t *)AscendC::GmAlloc(outputByteSize2);

    ReadFile("./input/input_x.bin", inputByteSize, x, inputByteSize);
    ReadFile("./input/input_y.bin", inputByteSize, y, inputByteSize);
    ReadFile("./input/input_x.bin", inputByteSize, x, inputByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(l1lossgrad_custom, blockDim, x, y, z, s ,u, reduction,*reinterpret_cast<L1LossGradCustomTilingData *>(tiling)); // use this macro for cpu debug

    WriteFile("./output/output_z.bin", z, outputByteSize);
    WriteFile("./output/output_s.bin", s, outputByteSize2);

    AscendC::GmFree((void *)x);
    AscendC::GmFree((void *)y);
    AscendC::GmFree((void *)z);
    AscendC::GmFree((void *)s);
    AscendC::GmFree((void *)tiling);
    AscendC::GmFree((void *)u);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    L1LossGradCustomTilingData *tiling;
    uint8_t *xHost, *yHost, *zHost,*sHost,*uHost;
    uint8_t *xDevice, *yDevice, *zDevice,*sDevice,*uDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize));
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);

    CHECK_ACL(aclrtMallocHost((void **)(&xHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&yHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&zHost), outputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&sHost), outputByteSize2));
    CHECK_ACL(aclrtMallocHost((void **)(&uHost), inputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&uDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&zDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&sDevice, outputByteSize2, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_x.bin", inputByteSize, xHost, inputByteSize);
    ReadFile("./input/input_y.bin", inputByteSize, yHost, inputByteSize);
    ReadFile("./input/input_s.bin", outputByteSize2, sHost, outputByteSize2);
    ReadFile("./input/grad_output_none.bin", inputByteSize, uHost, inputByteSize);


    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(yDevice, inputByteSize, yHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(uDevice, inputByteSize, uHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    
    ACLRT_LAUNCH_KERNEL(l1lossgrad_custom)(blockDim, stream, xDevice, yDevice, zDevice,sDevice,uDevice,reduction,tiling);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(zHost, outputByteSize, zDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, outputByteSize);

     CHECK_ACL(aclrtMemcpy(sHost, outputByteSize2, sDevice, outputByteSize2, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_s.bin", sHost, outputByteSize2);

    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFree(sDevice));
    CHECK_ACL(aclrtFree(uDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtFreeHost(zHost));
    CHECK_ACL(aclrtFreeHost(sHost));
    CHECK_ACL(aclrtFreeHost(sHost));
    CHECK_ACL(aclrtFreeHost(tiling));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}
